from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class OllamaConfig(BaseSettings):
    # 去除主观内容时使用
    OLLAMA_LLM_MODEL: str = Field(
        description="LLM model",
        default="llama3.1:70b",
        # default="llama3.1:8b",
        # default="llama3.1:70b-instruct-q2_K",
    )
    OLLAMA_LLM_URL: str = Field(
        description="LLM API URL",
        default="http://10.40.4.158:11434",
    )
    TEMPERATURE: float = Field(
        description="LLM temperature parameter",
        default=0.0,
    )

    OLLAMA_EMBEDDING_URL: str = Field(
        description="Embedding API URL",
        default="http://10.40.4.158:11436",
    )

    OLLAMA_EMBEDDING_MODEL: str = Field(
        description="Embedding model",
        default="nomic-embed-text",
    )


class DocumentConfig(BaseSettings):
    DOCUMENT_MD_DIR: str = Field(
        description="Document path",
        default="input/md",
    )


class SplitConfig(BaseSettings):
    CHUNK_SIZE: int = Field(
        description="Average chunk size",
        default=300,
    )
    OVERLAP_SIZE: int = Field(
        description="Overlap size",
        default=30,
    )
    CACHE: bool = Field(
        description="Cache",
        default=True,
    )


class LogConfig(BaseSettings):
    LOG_LEVEL: str = Field(
        description="Log level",
        default="INFO",
    )


class RAGConfig(BaseSettings):
    RERANCK_TOP_K: int = Field(
        description="Rerank top K, the number of documents to rerank",
        default=7,
    )


class HuggingFaceConfig(BaseSettings):
    RERANK_MODEL: str = Field(
        description="Rerank model",
        default="BAAI/bge-reranker-large",
    )


class AppConfig(
    OllamaConfig, LogConfig, SplitConfig, RAGConfig, DocumentConfig, HuggingFaceConfig
):
    model_config = SettingsConfigDict(
        # read from dotenv format config file
        env_file=".env",
        env_file_encoding="utf-8",
        frozen=True,
        # ignore extra attributes
        extra="ignore",
    )
    ANSWER_SHEET_PATH: str = Field(
        default="input/platform-policy-compliance/answer_sheet.csv"
    )
